package org.springframework.cloud.gateway.handler.predicate;

import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.gateway.handler.AsyncPredicate;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * @description: 重写此类，此类有bug，内存没有被回收
 * @author: Hou Dayu
 * @date: Created in 2020/11/19
 */
public class ReadBodyPredicateFactory extends AbstractRoutePredicateFactory<ReadBodyPredicateFactory.Config> {
    protected static final Log LOGGER = LogFactory.getLog(ReadBodyPredicateFactory.class);
    private static final String TEST_ATTRIBUTE = "read_body_predicate_test_attribute";
    private static final String CACHE_REQUEST_BODY_OBJECT_KEY = "cachedRequestBodyObject";
    private static final List<HttpMessageReader<?>> messageReaders = HandlerStrategies.withDefaults().messageReaders();

    public ReadBodyPredicateFactory() {
        super(ReadBodyPredicateFactory.Config.class);
    }

    public AsyncPredicate<ServerWebExchange> applyAsync(ReadBodyPredicateFactory.Config config) {
        return (exchange) -> {
            Class inClass = config.getInClass();
            Object cachedBody = exchange.getAttribute(CACHE_REQUEST_BODY_OBJECT_KEY);
            if (cachedBody != null) {
                try {
                    boolean test = config.predicate.test(cachedBody);
                    exchange.getAttributes().put(TEST_ATTRIBUTE, test);
                    return Mono.just(test);
                } catch (ClassCastException var7) {
                    if (LOGGER.isDebugEnabled()) {
                        LOGGER.debug("Predicate test failed because class in predicate does not match the cached body object", var7);
                    }

                    return Mono.just(false);
                }
            } else {
                return DataBufferUtils.join(exchange.getRequest().getBody()).flatMap((dataBuffer) -> {
                    DataBufferUtils.retain(dataBuffer);
                    final Flux<DataBuffer> cachedFlux = Flux.defer(() -> {
                        return Flux.just(dataBuffer.slice(0, dataBuffer.readableByteCount()));
                    });
                    ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
                        public Flux<DataBuffer> getBody() {
                            return cachedFlux;
                        }
                    };
                    ///////新增如下代码////////////////
                    DataBufferUtils.release(dataBuffer);
                    return ServerRequest.create(exchange.mutate().request(mutatedRequest).build(), messageReaders).bodyToMono(inClass).doOnNext((objectValue) -> {
                        exchange.getAttributes().put(CACHE_REQUEST_BODY_OBJECT_KEY, objectValue);
                        exchange.getAttributes().put("cachedRequestBody", cachedFlux);
                    }).map((objectValue) -> {
                        return config.predicate.test(objectValue);
                    });
                });
            }
        };
    }

    public Predicate<ServerWebExchange> apply(ReadBodyPredicateFactory.Config config) {
        throw new UnsupportedOperationException("ReadBodyPredicateFactory is only async.");
    }

    public static class Config {
        private Class inClass;
        private Predicate predicate;
        private Map<String, Object> hints;

        public Config() {
        }

        public Class getInClass() {
            return this.inClass;
        }

        public ReadBodyPredicateFactory.Config setInClass(Class inClass) {
            this.inClass = inClass;
            return this;
        }

        public Predicate getPredicate() {
            return this.predicate;
        }

        public <T> ReadBodyPredicateFactory.Config setPredicate(Class<T> inClass, Predicate<T> predicate) {
            this.setInClass(inClass);
            this.predicate = predicate;
            return this;
        }

        public ReadBodyPredicateFactory.Config setPredicate(Predicate predicate) {
            this.predicate = predicate;
            return this;
        }

        public Map<String, Object> getHints() {
            return this.hints;
        }

        public ReadBodyPredicateFactory.Config setHints(Map<String, Object> hints) {
            this.hints = hints;
            return this;
        }
    }
}